{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Terminal.py 功能测试\n",
    "\n",
    "本notebook用于测试terminal.py中的核心功能，包括：\n",
    "1. WebSocket连接建立\n",
    "2. 终端会话创建\n",
    "3. 数据发送和接收\n",
    "4. 心跳机制\n",
    "5. 终端大小调整"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. 测试连接建立...\n",
      "收到初始响应: \n",
      "The default interactive shell is now zsh.\n",
      "To update your account to use zsh, please run `chsh -s /bin/zsh`.\n",
      "For more details, please visit https://support.apple.com/kb/HT208050.\n",
      "\n",
      "\n",
      "2. 测试终端大小调整...\n",
      "\u001b[?1034h\u001b]633;A\u0007(byzerllm) (byzerllm) [w@me auto-coder.web]$ \u001b]633;B\u0007\n",
      "\u001b[K\u001b]633;A\u0007(byzerllm) (byzerllm) [w@me auto-coder.web]$ \u001b]633;B\u000733;B\u0007\n",
      "\n",
      "3. 测试命令执行...\n",
      "echo 'Hello Terminal'\n",
      "Hello Terminal\n",
      "\u001b]633;A\u0007(byzerllm) (byzerllm) [w@me auto-coder.web]$ \u001b]633;B\u0007\n",
      "命令执行响应: echo 'Hello Terminal'\n",
      "Hello Terminal\n",
      "\u001b]633;A\u0007(byzerllm) (byzerllm) [w@me auto-coder.web]$ \u001b]633;B\u0007\n",
      "\n",
      "所有测试通过！\n"
     ]
    }
   ],
   "source": [
    "import websockets\n",
    "import asyncio\n",
    "import json\n",
    "import time\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "@asynccontextmanager\n",
    "async def terminal_connection():\n",
    "    \"\"\"创建到terminal WebSocket的连接\"\"\"\n",
    "    uri = \"ws://localhost:8007/ws/terminal\"\n",
    "    async with websockets.connect(uri) as websocket:\n",
    "        yield websocket\n",
    "\n",
    "async def test_terminal_basic_functionality():\n",
    "    \"\"\"测试终端基本功能\"\"\"\n",
    "    async with terminal_connection() as ws:\n",
    "        # 测试1: 验证连接建立\n",
    "        print(\"1. 测试连接建立...\")\n",
    "        initial_response = await ws.recv()\n",
    "        print(f\"收到初始响应: {initial_response}\")        \n",
    "        \n",
    "        # 测试2: 发送终端调整大小命令\n",
    "        print(\"\\n2. 测试终端大小调整...\")\n",
    "        resize_data = {\n",
    "            \"type\": \"resize\",\n",
    "            \"cols\": 80,\n",
    "            \"rows\": 24\n",
    "        }\n",
    "        await ws.send(json.dumps(resize_data))\n",
    "\n",
    "        # 等待并收集输出\n",
    "        response = \"\"\n",
    "        start_time = time.time()\n",
    "        while time.time() - start_time < 2:  # 等待2秒\n",
    "            try:\n",
    "                msg = await asyncio.wait_for(ws.recv(), timeout=0.5)\n",
    "                response += msg if isinstance(msg, str) else msg.decode('utf-8')\n",
    "                print(response)\n",
    "            except asyncio.TimeoutError:\n",
    "                continue\n",
    "        \n",
    "        # 测试3: 发送简单命令并验证响应\n",
    "        print(\"\\n3. 测试命令执行...\")\n",
    "        await ws.send(\"echo 'Hello Terminal'\\n\")\n",
    "        \n",
    "        # 等待并收集输出\n",
    "        response = \"\"\n",
    "        start_time = time.time()\n",
    "        while time.time() - start_time < 10:  # 等待2秒\n",
    "            try:\n",
    "                msg = await asyncio.wait_for(ws.recv(), timeout=0.5)\n",
    "                response += msg if isinstance(msg, str) else msg.decode('utf-8')\n",
    "                print(response)\n",
    "                if 'Hello Terminal' in response:\n",
    "                    break\n",
    "            except asyncio.TimeoutError:\n",
    "                continue\n",
    "                \n",
    "        print(f\"命令执行响应: {response}\")\n",
    "        assert 'Hello Terminal' in response                        \n",
    "        \n",
    "        print(\"\\n所有测试通过！\")\n",
    "\n",
    "# 运行测试\n",
    "await test_terminal_basic_functionality()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 压力测试\n",
    "\n",
    "测试多个并发连接的情况"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def stress_test(num_connections=5):\n",
    "    \"\"\"创建多个并发连接进行压力测试\"\"\"\n",
    "    print(f\"开始{num_connections}个并发连接的压力测试...\")\n",
    "    \n",
    "    async def single_connection_test(conn_id):\n",
    "        try:\n",
    "            async with terminal_connection() as ws:\n",
    "                # 发送一些基本命令\n",
    "                await ws.send(f\"echo 'Connection {conn_id}'\\n\")\n",
    "                response = await ws.recv()\n",
    "                print(f\"连接 {conn_id} 收到响应: {response}\")\n",
    "                \n",
    "                # 保持连接一段时间\n",
    "                await asyncio.sleep(2)\n",
    "                \n",
    "                return True\n",
    "        except Exception as e:\n",
    "            print(f\"连接 {conn_id} 发生错误: {str(e)}\")\n",
    "            return False\n",
    "    \n",
    "    # 创建多个并发连接\n",
    "    tasks = [single_connection_test(i) for i in range(num_connections)]\n",
    "    results = await asyncio.gather(*tasks)\n",
    "    \n",
    "    # 统计结果\n",
    "    success_count = sum(1 for r in results if r)\n",
    "    print(f\"\\n压力测试完成:\")\n",
    "    print(f\"成功连接: {success_count}/{num_connections}\")\n",
    "    \n",
    "# 运行压力测试\n",
    "await stress_test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"# Docker Push Guide\\n\\n## byzer-storage\\n\\nThis guide explains how to use the byzer-storage Docker image.\\n\\n### Image Information\\n\\nThe byzer-storage Docker image is based on Ubuntu 22.04 and includes:\\n- Miniconda with Python 3.10.11\\n- auto-coder package\\n- python-multipart package\\n\\n### Usage\\n\\nTo use this image:\\n\\n```bash\\n# Pull the image\\ndocker pull [repository]/byzer-storage:latest\\n\\n# Run a container\\ndocker run -it [repository]/byzer-storage:latest\\n```\\n\\n### Building the Image\\n\\nTo build the image locally:\\n\\n```bash\\ncd /path/to/auto-coder.web\\ndocker build -f docker/base/Dockerfile -t byzer-storage:latest .\\n```\\n\\n### Configuration\\n\\nThe image is configured with:\\n- Python 3.10.11 in a conda environment named 'py310'\\n- PIP configured to use Tsinghua mirror\\n- auto-coder and python-multipart packages pre-installed\""
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from autocoder.common import files as FileUtils\n",
    "f = \"/Users/allwefantasy/projects/auto-coder.web/docker-push-guide.md\"\n",
    "FileUtils.read_file(f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
